package cn.infomany.code;

import cn.infomany.code.generate.CodeBaseVisitor;
import cn.infomany.code.generate.CodeLexer;
import cn.infomany.code.generate.CodeParser;

import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @description: 范围访问器
 * @author: zhanjinbing
 * @data: 2020-03-03 15:17
 */
public class RangesVisitor extends CodeBaseVisitor<Object> {

    @Override
    public Object visitAdd(CodeParser.AddContext ctx) {
        Set<Integer> result = new HashSet<>();
        Set<Integer> result0 = (Set<Integer>) visit(ctx.range(0));
        Set<Integer> result1 = (Set<Integer>) visit(ctx.range(1));

        result.addAll(result0);
        result.addAll(result1);
        return result;
    }

    @Override
    public Object visitSE(CodeParser.SEContext ctx) {
        int speed = 1;
        int start = (int) visit(ctx.start());
        int end = (int) visit(ctx.end());

        return visitSSE(start, speed, end,
                ctx.start.getType() == CodeLexer.MEDIUMBRACKETLEFT,
                ctx.stop.getType() == CodeLexer.MEDIUMBRACKETRIGHT);

    }

    private Set<Integer> visitSSE(int start, int speed, int end, boolean hasStart, boolean hasEnd) {
        Set<Integer> result = new HashSet<>();
        for (int i = start; i < end; i += speed) {
            result.add(i);
        }

        if (!hasStart) {
            result.remove(start);
        }

        if (hasEnd) {
            result.add(end);
        }

        return result;
    }

    @Override
    public Object visitSimple(CodeParser.SimpleContext ctx) {
        return visit(ctx.range());
    }

    @Override
    public Object visitS(CodeParser.SContext ctx) {
        Set<Integer> results = new HashSet<>();
        if (ctx.start.getType() == CodeLexer.SMALLBRACKETLEFT
                || ctx.stop.getType() == CodeLexer.SMALLBRACKETRIGHT) {
            return results;
        }

        results.add((Integer) visit(ctx.start()));
        return results;
    }

    @Override
    public Object visitSSE(CodeParser.SSEContext ctx) {
        Object speedObj = visit(ctx.speed());
        int start = (int) visit(ctx.start());
        int end = (int) visit(ctx.end());

        Set<Integer> results;

        if (speedObj instanceof String) {
            results = visitSSE(start, 1, end,
                    ctx.start.getType() == CodeLexer.MEDIUMBRACKETLEFT,
                    ctx.stop.getType() == CodeLexer.MEDIUMBRACKETRIGHT);

            final Integer tmpInt = Integer.valueOf(((String) speedObj).substring(1));

            results = results.stream().filter(it -> it % tmpInt == 0).collect(Collectors.toSet());

        } else if (speedObj instanceof Integer) {
            results = visitSSE(start, (Integer) speedObj, end,
                    ctx.start.getType() == CodeLexer.MEDIUMBRACKETLEFT,
                    ctx.stop.getType() == CodeLexer.MEDIUMBRACKETRIGHT);
        } else {
            throw new RuntimeException("[speed]格式有误");
        }
        return results;
    }

    @Override
    public Object visitAssoc(CodeParser.AssocContext ctx) {
        Set<Integer> result = new HashSet<>();

        Set<Integer> result0 = (Set<Integer>) visit(ctx.range(0));
        Set<Integer> result1 = (Set<Integer>) visit(ctx.range(1));

        if (ctx.op().MixOPTION() != null) {
            result0.forEach(it -> {
                if (result1.contains(it)) {
                    result.add(it);
                }
            });
        } else {
            result.addAll(result0);
            if (ctx.op().ADDOPTION() != null) {
                result.addAll(result1);
            } else {
                result.removeAll(result1);
            }
        }
        return result;
    }

    @Override
    public Object visitStart(CodeParser.StartContext ctx) {
        return Integer.valueOf(ctx.INT().getText());
    }

    @Override
    public Object visitSpeed(CodeParser.SpeedContext ctx) {
        //  检查是不是单纯的int
        String reg = "\\d+";
        if (ctx.getText().matches(reg)) {
            return Integer.valueOf(ctx.INT().getText());
        }
        return ctx.getText();
    }

    @Override
    public Object visitEnd(CodeParser.EndContext ctx) {
        return Integer.valueOf(ctx.INT().getText());
    }
}